package club.janna.springboot.study.interceptor;

import club.janna.springboot.study.util.JannaDataSourceHolder;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.tomcat.util.collections.ManagedConcurrentWeakHashMap;

import java.util.Map;
import java.util.WeakHashMap;

@Slf4j
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class})
})
public class JannaRouteDataSourceInterceptor implements Interceptor {

    private final Map<Object, JannaDataSourceHolder.Type> cacheMap = new ManagedConcurrentWeakHashMap<>();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        JannaDataSourceHolder.Type t;
        if(!cacheMap.containsKey(ms)) {
            SqlCommandType sqlCommandType = ms.getSqlCommandType();
            log.info("sql type -> {}", sqlCommandType);
            t = sqlCommandType == SqlCommandType.SELECT ? JannaDataSourceHolder.Type.READ : JannaDataSourceHolder.Type.WRITE;
            log.info("type -> {}", t);
            cacheMap.put(ms, t);
        } else
            t = cacheMap.get(ms);
        JannaDataSourceHolder.set(t);
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return target instanceof Executor ? Plugin.wrap(target, this) : target;
    }
}
